"""
并发控制机制
使用Redis实现分布式信号量和速率限制
"""
import redis
import time
import logging
from contextlib import contextmanager
from typing import Optional
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

class ConcurrencyController:
    """并发控制器"""
    
    def __init__(self, redis_host='localhost', redis_port=6379, redis_db=2):
        """
        初始化并发控制器
        
        Args:
            redis_host: Redis主机
            redis_port: Redis端口
            redis_db: Redis数据库号
        """
        self.redis_client = redis.Redis(
            host=redis_host,
            port=redis_port,
            db=redis_db,
            decode_responses=True
        )
        
        # 平台并发限制配置
        self.platform_limits = {
            'weibo': 3,
            'zhihu': 2,
            'toutiao': 3,
            'baidu': 3,
            'xiaohongshu': 2,
            'douyin': 2,
            'twitter': 2,
            'reddit': 3,
            'youtube': 1,
        }
        
        # 速率限制配置（每分钟最大请求数）
        self.rate_limits = {
            'weibo': 10,
            'zhihu': 15,
            'toutiao': 20,
            'baidu': 20,
            'xiaohongshu': 10,
            'douyin': 10,
            'twitter': 15,
            'reddit': 20,
            'youtube': 10,
        }
        
    @contextmanager
    def acquire_semaphore(self, platform: str, timeout: int = 30):
        """
        获取平台信号量
        
        Args:
            platform: 平台名称
            timeout: 获取超时时间（秒）
            
        Yields:
            是否成功获取
        """
        semaphore_key = f'semaphore:{platform}'
        limit = self.platform_limits.get(platform, 1)
        acquired = False
        start_time = time.time()
        
        try:
            # 尝试获取信号量
            while time.time() - start_time < timeout:
                current = self.redis_client.get(semaphore_key)
                current_count = int(current) if current else 0
                
                if current_count < limit:
                    # 使用Lua脚本保证原子性
                    lua_script = """
                    local key = KEYS[1]
                    local limit = tonumber(ARGV[1])
                    local current = redis.call('GET', key)
                    local count = current and tonumber(current) or 0
                    
                    if count < limit then
                        redis.call('INCR', key)
                        redis.call('EXPIRE', key, 3600)
                        return 1
                    else
                        return 0
                    end
                    """
                    
                    result = self.redis_client.eval(lua_script, 1, semaphore_key, limit)
                    if result == 1:
                        acquired = True
                        logger.info(f"Acquired semaphore for {platform}")
                        break
                        
                time.sleep(1)
                
            if not acquired:
                logger.warning(f"Failed to acquire semaphore for {platform} within {timeout}s")
                
            yield acquired
            
        finally:
            if acquired:
                # 释放信号量
                self.redis_client.decr(semaphore_key)
                logger.info(f"Released semaphore for {platform}")
                
    def check_rate_limit(self, platform: str) -> bool:
        """
        检查速率限制
        
        Args:
            platform: 平台名称
            
        Returns:
            是否允许执行
        """
        rate_key = f'rate:{platform}:{int(time.time() / 60)}'
        limit = self.rate_limits.get(platform, 10)
        
        # 使用Lua脚本实现滑动窗口
        lua_script = """
        local key = KEYS[1]
        local limit = tonumber(ARGV[1])
        local current = redis.call('GET', key)
        local count = current and tonumber(current) or 0
        
        if count < limit then
            redis.call('INCR', key)
            redis.call('EXPIRE', key, 60)
            return 1
        else
            return 0
        end
        """
        
        result = self.redis_client.eval(lua_script, 1, rate_key, limit)
        
        if result == 1:
            logger.debug(f"Rate limit check passed for {platform}")
            return True
        else:
            logger.warning(f"Rate limit exceeded for {platform}")
            return False
            
    def get_queue_length(self, queue_name: str) -> int:
        """
        获取队列长度
        
        Args:
            queue_name: 队列名称
            
        Returns:
            队列长度
        """
        # Celery使用的队列key格式
        queue_key = f'celery:{queue_name}'
        return self.redis_client.llen(queue_key)
        
    def set_max_queue_length(self, queue_name: str, max_length: int):
        """
        设置队列最大长度
        
        Args:
            queue_name: 队列名称
            max_length: 最大长度
        """
        config_key = f'queue_config:{queue_name}:max_length'
        self.redis_client.set(config_key, max_length)
        
    def check_queue_limit(self, queue_name: str) -> bool:
        """
        检查队列是否超过限制
        
        Args:
            queue_name: 队列名称
            
        Returns:
            是否可以添加任务
        """
        config_key = f'queue_config:{queue_name}:max_length'
        max_length = self.redis_client.get(config_key)
        
        if not max_length:
            return True  # 没有限制
            
        current_length = self.get_queue_length(queue_name)
        return current_length < int(max_length)
        
    def get_platform_stats(self, platform: str) -> dict:
        """
        获取平台统计信息
        
        Args:
            platform: 平台名称
            
        Returns:
            统计信息
        """
        semaphore_key = f'semaphore:{platform}'
        rate_key = f'rate:{platform}:{int(time.time() / 60)}'
        
        current_semaphore = self.redis_client.get(semaphore_key)
        current_rate = self.redis_client.get(rate_key)
        
        return {
            'platform': platform,
            'current_concurrent': int(current_semaphore) if current_semaphore else 0,
            'max_concurrent': self.platform_limits.get(platform, 1),
            'current_rate': int(current_rate) if current_rate else 0,
            'max_rate': self.rate_limits.get(platform, 10),
            'timestamp': datetime.utcnow().isoformat()
        }
        
    def reset_platform_limits(self, platform: str):
        """
        重置平台限制
        
        Args:
            platform: 平台名称
        """
        semaphore_key = f'semaphore:{platform}'
        self.redis_client.delete(semaphore_key)
        
        # 清理速率限制键
        pattern = f'rate:{platform}:*'
        for key in self.redis_client.scan_iter(match=pattern):
            self.redis_client.delete(key)
            
        logger.info(f"Reset limits for platform: {platform}")
        
    def update_platform_limit(self, platform: str, concurrent_limit: Optional[int] = None,
                            rate_limit: Optional[int] = None):
        """
        更新平台限制
        
        Args:
            platform: 平台名称
            concurrent_limit: 并发限制
            rate_limit: 速率限制
        """
        if concurrent_limit is not None:
            self.platform_limits[platform] = concurrent_limit
            logger.info(f"Updated concurrent limit for {platform}: {concurrent_limit}")
            
        if rate_limit is not None:
            self.rate_limits[platform] = rate_limit
            logger.info(f"Updated rate limit for {platform}: {rate_limit}")

# 创建全局控制器实例
concurrency_controller = ConcurrencyController()